#if USE_NUMA == 1

#include <cassert>

#include <hip/hip_runtime.h>

#include "scope/init/flags.hpp"
#include "scope/init/init.hpp"
#include "scope/utils/utils.hpp"

#include "args.hpp"
#include "init/flags.hpp"
#include "init/numa.hpp"
#include "utils/numa.hpp"

#define NAME "Comm/NUMAMemcpy/GPUToGPU"

static void Comm_NUMAMemcpy_GPUToGPU(benchmark::State &state) {

  if (!has_cuda) {
    state.SkipWithError(NAME " no CUDA device found");
    return;
  }

  if (!has_numa) {
    state.SkipWithError(NAME " no NUMA control available");
    return;
  }

  if (num_gpus() < 2) {
    state.SkipWithError(NAME " requires at least two GPUs");
    return;
  }
  
  const int numa_id = FLAG(numa_ids)[0];
  const int src_gpu = FLAG(cuda_device_ids)[0];
  const int dst_gpu = FLAG(cuda_device_ids)[1];

  if (src_gpu == dst_gpu) {
    state.SkipWithError(NAME " requires two different GPUs");
    return;
  }

  const auto bytes  = 1ULL << static_cast<size_t>(state.range(0));


  numa_bind_node(numa_id);

  char *src = nullptr;
  char *dst = nullptr;

  if (PRINT_IF_ERROR(utils::cuda_reset_device(src_gpu))) {
    state.SkipWithError(NAME " failed to reset CUDA device");
    return;
  }
  if (PRINT_IF_ERROR(utils::cuda_reset_device(dst_gpu))) {
    state.SkipWithError(NAME " failed to reset CUDA device");
    return;
  }

  if (PRINT_IF_ERROR(hipSetDevice(src_gpu))) {
    state.SkipWithError(NAME " failed to set src device");
    return;
  }
  if (PRINT_IF_ERROR(hipMalloc(&src, bytes))) {
    state.SkipWithError(NAME " failed to perform hipMalloc");
    return;
  }
  defer(hipFree(src));
  if (PRINT_IF_ERROR(hipMemset(src, 0, bytes))) {
    state.SkipWithError(NAME " failed to perform src hipMemset");
    return;
  }
  hipError_t err = hipDeviceDisablePeerAccess(dst_gpu);
  if (hipSuccess != err && hipErrorPeerAccessNotEnabled != err) {
    state.SkipWithError(NAME " failed to disable peer access");
    return;
  }

  if (PRINT_IF_ERROR(hipSetDevice(dst_gpu))) {
    state.SkipWithError(NAME " failed to set dst device");
    return;
  }
  if (PRINT_IF_ERROR(hipMalloc(&dst, bytes))) {
    state.SkipWithError(NAME " failed to perform hipMalloc");
    return;
  }
  defer(hipFree(dst));
  if (PRINT_IF_ERROR(hipMemset(dst, 0, bytes))) {
    state.SkipWithError(NAME " failed to perform dst hipMemset");
    return;
  }
  err = hipDeviceDisablePeerAccess(src_gpu);
  if (hipSuccess != err && hipErrorPeerAccessNotEnabled != err) {
    state.SkipWithError(NAME " failed to disable peer access");
    return;
  }

  hipEvent_t start, stop;
  PRINT_IF_ERROR(hipEventCreate(&start));
  PRINT_IF_ERROR(hipEventCreate(&stop));

  for (auto _ : state) {

    hipEventRecord(start, NULL);
    const auto cuda_err = hipMemcpy(dst, src, bytes, hipMemcpyDeviceToDevice);
    hipEventRecord(stop, NULL);
    hipEventSynchronize(stop);

    if (PRINT_IF_ERROR(cuda_err) != hipSuccess) {
      state.SkipWithError(NAME " failed to perform memcpy");
      break;
    }
    float msecTotal = 0.0f;
    if (PRINT_IF_ERROR(hipEventElapsedTime(&msecTotal, start, stop))) {
      state.SkipWithError(NAME " failed to get elapsed time");
      break;
    }
    state.SetIterationTime(msecTotal / 1000);
  }
  state.SetBytesProcessed(int64_t(state.iterations()) * int64_t(bytes));
  state.counters.insert({{"bytes", bytes}});

  // re-enable NUMA scheduling
  numa_bind_node(-1);

  // re-enable peer access
  err = hipSetDevice(src_gpu);
  err = hipDeviceEnablePeerAccess(dst_gpu, 0);
  // docs say should return hipErrorInvalidDevice, actually returns hipErrorPeerAccessUnsupported?
  if (hipSuccess != err && hipErrorInvalidDevice != err && hipErrorPeerAccessUnsupported != err) {
    PRINT_IF_ERROR(err);
    state.SkipWithError(NAME " couldn't re-enable peer access");
  }
  err = hipSetDevice(dst_gpu);
  err = hipDeviceEnablePeerAccess(src_gpu, 0);
  if (hipSuccess != err && hipErrorInvalidDevice != err && hipErrorPeerAccessUnsupported != err) {
    PRINT_IF_ERROR(err);
    state.SkipWithError(NAME " couldn't re-enable peer access");
  }

  numa_bind_node(-1);
}

BENCHMARK(Comm_NUMAMemcpy_GPUToGPU)->SMALL_ARGS()->UseManualTime();

#endif // USE_NUMA == 1